"""交通场景分析器模块"""

import cv2
import numpy as np
from typing import Dict, List, Tuple, Optional
import time

class TrafficAnalyzer:
    """交通场景分析器，负责处理交通相关的分析任务"""
    
    def __init__(self):
        """初始化交通分析器"""
        # 交通类别定义
        self.traffic_classes = {
            'person': {'color': (0, 255, 0), 'label': '行人'},
            'car': {'color': (255, 0, 0), 'label': '小汽车'},
            'bicycle': {'color': (0, 255, 255), 'label': '自行车'},
            'motorbike': {'color': (255, 255, 0), 'label': '摩托车'},
            'truck': {'color': (255, 0, 255), 'label': '卡车'},
            'bus': {'color': (0, 0, 255), 'label': '公交车'}
        }
        
        # 统计信息
        self.reset_statistics()
        
        # 历史轨迹
        self.tracks = {}
        
        print("交通分析器已初始化")
    
    def reset_statistics(self):
        """重置统计信息"""
        self.statistics = {
            'total_detected': 0,
            'by_class': {cls: 0 for cls in self.traffic_classes},
            'speeds': [],
            'avg_speeds': {cls: [] for cls in self.traffic_classes},
            'last_update_time': time.time()
        }
    
    def analyze_traffic(self, frame: np.ndarray, detections: List[Dict], 
                       tracking_results: List[Dict] = None) -> Dict:
        """分析交通场景
        
        Args:
            frame: 输入视频帧
            detections: YOLO检测结果
            tracking_results: 跟踪结果（可选）
            
        Returns:
            交通分析结果字典
        """
        results = {
            'detections_count': {cls: 0 for cls in self.traffic_classes},
            'traffic_density': 0,
            'moving_objects': [],
            'queue_lengths': {},
            'annotations': []
        }
        
        # 统计各类别数量
        for detection in detections:
            cls = detection.get('class', 'unknown')
            if cls in self.statistics['by_class']:
                results['detections_count'][cls] += 1
                self.statistics['by_class'][cls] += 1
                self.statistics['total_detected'] += 1
        
        # 计算交通密度
        height, width = frame.shape[:2]
        total_pixels = height * width
        occupied_area = 0
        
        for detection in detections:
            bbox = detection.get('bbox', [])
            if len(bbox) >= 4:
                # 简化的面积计算
                x1, y1, x2, y2 = bbox
                occupied_area += (x2 - x1) * (y2 - y1)
        
        # 交通密度 = 被检测对象占据的像素比例
        results['traffic_density'] = occupied_area / total_pixels if total_pixels > 0 else 0
        
        # 处理跟踪结果
        if tracking_results:
            moving_objects = []
            for track in tracking_results:
                # 这里简化处理，实际项目中可以计算速度、方向等
                if 'speed' in track and track['speed'] > 0:
                    moving_objects.append({
                        'id': track.get('id'),
                        'class': track.get('class'),
                        'speed': track.get('speed'),
                        'direction': track.get('direction', 'unknown')
                    })
            
            results['moving_objects'] = moving_objects
            
            # 更新速度统计
            for obj in moving_objects:
                cls = obj.get('class')
                speed = obj.get('speed')
                if cls in self.statistics['avg_speeds'] and speed > 0:
                    self.statistics['avg_speeds'][cls].append(speed)
                    self.statistics['speeds'].append(speed)
        
        # 生成注释信息
        annotations = []
        for detection in detections:
            bbox = detection.get('bbox', [])
            cls = detection.get('class', 'unknown')
            confidence = detection.get('confidence', 0)
            
            if cls in self.traffic_classes and len(bbox) >= 4:
                annotations.append({
                    'bbox': bbox,
                    'class': cls,
                    'label': self.traffic_classes[cls]['label'],
                    'color': self.traffic_classes[cls]['color'],
                    'confidence': confidence
                })
        
        results['annotations'] = annotations
        
        # 更新最后统计时间
        self.statistics['last_update_time'] = time.time()
        
        return results
    
    def get_statistics(self) -> Dict:
        """获取当前统计信息
        
        Returns:
            统计信息字典
        """
        # 计算平均速度
        avg_speeds = {}
        for cls, speeds in self.statistics['avg_speeds'].items():
            if speeds:
                avg_speeds[cls] = sum(speeds) / len(speeds)
            else:
                avg_speeds[cls] = 0
        
        overall_avg_speed = 0
        if self.statistics['speeds']:
            overall_avg_speed = sum(self.statistics['speeds']) / len(self.statistics['speeds'])
        
        return {
            'total_detected': self.statistics['total_detected'],
            'by_class': self.statistics['by_class'],
            'average_speeds': avg_speeds,
            'overall_average_speed': overall_avg_speed,
            'last_update_time': self.statistics['last_update_time']
        }
    
    def visualize_analysis(self, frame: np.ndarray, analysis_results: Dict) -> np.ndarray:
        """可视化分析结果
        
        Args:
            frame: 输入视频帧
            analysis_results: 分析结果
            
        Returns:
            带注释的视频帧
        """
        annotated_frame = frame.copy()
        
        # 绘制检测框和标签
        for annotation in analysis_results.get('annotations', []):
            x1, y1, x2, y2 = map(int, annotation['bbox'])
            color = annotation['color']
            label = f"{annotation['label']} {annotation['confidence']:.2f}"
            
            # 绘制边界框
            cv2.rectangle(annotated_frame, (x1, y1), (x2, y2), color, 2)
            
            # 绘制标签
            label_size, _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 2)
            label_y1 = max(y1 - 10, 0)
            cv2.rectangle(annotated_frame, 
                         (x1, label_y1 - label_size[1] - 10),
                         (x1 + label_size[0], label_y1),
                         color, -1)
            cv2.putText(annotated_frame, label, (x1, label_y1 - 5),
                       cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2)
        
        # 显示交通密度
        density_text = f"交通密度: {analysis_results.get('traffic_density', 0) * 100:.1f}%"
        cv2.putText(annotated_frame, density_text, (10, 30),
                   cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
        
        # 显示各类别数量
        y_offset = 60
        for cls, count in analysis_results.get('detections_count', {}).items():
            if count > 0 and cls in self.traffic_classes:
                count_text = f"{self.traffic_classes[cls]['label']}: {count}"
                cv2.putText(annotated_frame, count_text, (10, y_offset),
                           cv2.FONT_HERSHEY_SIMPLEX, 0.5, self.traffic_classes[cls]['color'], 2)
                y_offset += 25
        
        return annotated_frame